package com.ydlclass.loadbalance.loadbalancer;

import com.ydlclass.configuration.YrpcConfiguration;
import com.ydlclass.configuration.YrpcConfigurationHolder;
import com.ydlclass.loadbalance.LoadBalancer;
import com.ydlclass.loadbalance.Selector;
import com.ydlclass.message.MetaData;
import com.ydlclass.message.YrpcMessage;
import com.ydlclass.registry.Registry;
import lombok.extern.slf4j.Slf4j;

import java.net.InetSocketAddress;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 使用一致性hash算法进行负载均衡
 * @author it楠老师
 * @createTime 2023-03-26
 */
@Slf4j
public class ConsistentHashLoadBalancer implements LoadBalancer {

    // 针对每一个服务建立一个一致性hash的选择器
    private Map<String,ConsistentHashSelector> selectors = new ConcurrentHashMap<>(32);

    @Override
    public InetSocketAddress selectServiceAddress(List<InetSocketAddress> serviceInetSocketAddress, YrpcMessage yrpcMessage) {
        if(yrpcMessage.getPayload() instanceof MetaData metaData){
            String serviceName = metaData.getInterfaceName();
            // 根据服务名称获取对应的selector
            ConsistentHashSelector consistentHashSelector = selectors.get(serviceName);
            // 未命中则创建一个选择器
            if(consistentHashSelector == null){
                // 1、尝试重新拉取服务列表
                YrpcConfiguration yrpcConfiguration = YrpcConfigurationHolder.get();
                Registry registry = yrpcConfiguration.getRegistry();
                List<InetSocketAddress> serviceList = registry.lookup(serviceName);

                // 2、创建看一个一致性hash选择器
                consistentHashSelector = new ConsistentHashSelector(
                        serviceList,64
                );

                selectors.put(serviceName,consistentHashSelector);

            }
            // 为了简单起见，我们直接使用请求id进行hash路由，
            // 当然选择路由的方式很多比如客户端的ip等
            if(log.isDebugEnabled()){
                log.debug("balancing --> The consistent hash algorithm is being used for load balancing.");
            }
            return consistentHashSelector.getNext(yrpcMessage);
        }

        return null;
    }


    /**
     * 这是使用chatGpt生成的一个一致性hash算法
     */
    private static class ConsistentHashSelector implements Selector {
        private final SortedMap<Integer, InetSocketAddress> circle = new TreeMap<>(); // 存储哈希环上的节点
        private final int virtualNodes ; // 每个节点的虚拟节点数量

        // 构造函数，传入服务节点列表和每个节点的虚拟节点数量

        /**
         * @param nodes
         * @param virtualNodes
         */
        public ConsistentHashSelector(List<InetSocketAddress> nodes, int virtualNodes) {
            this.virtualNodes = virtualNodes;
            for (InetSocketAddress node : nodes) {
                addNode(node);
            }
        }

        // 添加服务节点和对应的虚拟节点到哈希环中
        private void addNode(InetSocketAddress node) {
            for (int i = 0; i < virtualNodes; i++) {
                int hash = hash(node + "-" + i); // 对服务节点和虚拟节点的字符串进行哈希，得到一个32位整数
                circle.put(hash, node); // 将哈希值和服务节点映射存储到哈希环上
            }
        }

        // 从哈希环上移除服务节点和对应的虚拟节点
        private void removeNode(InetSocketAddress node) {
            for (int i = 0; i < virtualNodes; i++) {
                int hash = hash(node + "-" + i); // 对服务节点和虚拟节点的字符串进行哈希，得到一个32位整数
                circle.remove(hash); // 从哈希环上移除对应的哈希值和服务节点映射
            }
        }

        // 根据请求的key值获取对应的服务节点
        @Override
        public InetSocketAddress getNext(YrpcMessage yrpcMessage) {
            String key = Long.toString(yrpcMessage.getRequestId());
            if (circle.isEmpty()) { // 如果哈希环为空，返回null
                return null;
            }

            int hash = hash(key); // 对请求的key值进行哈希，得到一个32位整数

            if (!circle.containsKey(hash)) { // 如果哈希环中不包含该哈希值，查找第一个大于该哈希值的节点
                SortedMap<Integer, InetSocketAddress> tailMap = circle.tailMap(hash);
                hash = tailMap.isEmpty() ? circle.firstKey() : tailMap.firstKey();
            }

            return circle.get(hash); // 返回对应的服务节点
        }

        // 根据字符串的哈希值获取一个32位整数
        private int hash(String key) {
            MessageDigest md;
            try {
                md = MessageDigest.getInstance("MD5");
            } catch (NoSuchAlgorithmException e) {
                throw new RuntimeException(e);
            }
            byte[] hash = md.digest(key.getBytes());
            // 将md5结果的后四个字节转化为int
            int result = 0;
            for(int i=0;i<4;i++){
                result <<= 8;//左移 8 位
                result |= hash[3-i];//保存 byte 值到 ans 的最低 8 位
            }
            return result;
        }

        @Override
        public void reload(String serviceName, List<InetSocketAddress> newServiceList) {

        }

    }


}
